# get colors for plots, total 63 colors can use
get_colors <- function(colors_len = NULL){
  cols <- c(RColorBrewer::brewer.pal(12,"Paired"),
           RColorBrewer::brewer.pal(9,"Set1"),
           RColorBrewer::brewer.pal(8,"Set2"),
           RColorBrewer::brewer.pal(12,"Set3"),
           RColorBrewer::brewer.pal(8,"Pastel2"),
           RColorBrewer::brewer.pal(9,"Pastel1"),
           RColorBrewer::brewer.pal(8,"Accent"))
  if(missing(colors_len)){
    col <- unique(cols)[-c(17,23)]
  }else{
    col <- unique(cols)[-c(17,23)][1:colors_len]
  }
  return(col)
}

# prep data to create monocle object
# 表达矩阵, 建议使用 count 矩阵
# 基因注释, 对基因进行注释, 行名为基因, 只要是关于基因的描述信息都可以, 如：基因对应的GO KEGG号等信息, 但是第一列必须为 "gene_short_name", 否则会报错, 所以避免麻烦我们就弄一列。
# 表型信息, 对细胞的描述文件, 行名为 cell barcode, 同样的 只要是关于细胞的描述信息都可以, 如：cluster, celltype, Sample ... , 没啥要求都可以加, 所以我们把 peoject@meta.data 全部提取出来给他。
# 最好是选取具有分化关系的 cluster 来进行 monocle2 轨迹分析
prep_data <- function(project = NULL,
                      cluster_type = NULL,
                      sub_cluster = NULL){
  if(!is.null(cluster_type) & !is.null(sub_cluster)){
    Idents(project) <- cluster_type
    project <- subset(project, idents = sub_cluster)
  }
  gene_count <- Seurat::GetAssayData(object = project[["RNA"]], slot = "count")
  gene_anno  <- data.frame(gene_id = rownames(gene_count), gene_short_name = rownames(gene_count)) %>% 
    tibble::column_to_rownames(var = "gene_id")
  pheno_data <- project@meta.data
  result <- list(count = gene_count, anno = gene_anno, pheno = pheno_data)
  return(result)
}

# plot_genes_branched_pseudotime 是不支持绘制单个基因的表达分布图的, 报错如下：
# Error in if (nrow(ancestor_exprs) == 1) exprs_data <- t(as.matrix(ancestor_exprs)) else exprs_data <- ancestor_exprs: argument is of length zero
# 所以不能传入单个基因, 自己根据 plot_genes_branched_pseudotime 源代码进行改造, 改造成 2 个 function, 一个 run_cala 和 run_plot
# gene_list <- c("Dkk3", "Ptn", "Lum", "Dcn", "Sparc", "Crabp1", "Spon2", "Tcf21", "Wif1")
# genes <- gene_list[gene_list %in% rownames(monocl_obj)]
# result <- run_calc(monocl_obj[genes,], branch_point = 1)
# run_plot(cds_exprs = result$data, min_expr = result$min_expr, genes = genes, name = "Fibroblasts", branch_point = 1)
monocle_theme_opts <- function(){
    theme(strip.background = element_rect(colour = 'white', fill = 'white')) +
    theme(panel.border = element_blank(), axis.line = element_line()) +
    theme(panel.grid.minor.x = element_blank(), panel.grid.minor.y = element_blank()) +
    theme(panel.grid.major.x = element_blank(), panel.grid.major.y = element_blank()) + 
    theme(panel.background = element_rect(fill='white')) +
    theme(legend.key = element_blank()) +
    theme(plot.title = element_text(size=12,hjust=0.5))
}

run_calc <- function (cds, branch_states = NULL, branch_point = 1, branch_labels = NULL, 
    method = "fitting", min_expr = NULL, cell_size = 0.75, nrow = NULL, 
    ncol = 1, panel_order = NULL, color_by = "State", expression_curve_linetype_by = "Branch", 
    trend_formula = "~ sm.ns(Pseudotime, df=3) * Branch", reducedModelFormulaStr = NULL, 
    label_by_short_name = TRUE, relative_expr = TRUE, ...){
    Branch <- NA
    if (is.null(reducedModelFormulaStr) == FALSE) {
        pval_df <- branchTest(cds, branch_states = branch_states, 
            branch_point = branch_point, fullModelFormulaStr = trend_formula, 
            reducedModelFormulaStr = "~ sm.ns(Pseudotime, df=3)", 
            ...)
        fData(cds)[, "pval"] <- pval_df[row.names(cds), "pval"]
    }
    if ("Branch" %in% all.vars(terms(as.formula(trend_formula)))) {
        cds_subset <- buildBranchCellDataSet(cds = cds, branch_states = branch_states, 
            branch_point = branch_point, branch_labels = branch_labels, 
            progenitor_method = "duplicate", ...)
    }
    else {
        cds_subset <- cds
        pData(cds_subset)$Branch <- pData(cds_subset)$State
    }
    if (cds_subset@expressionFamily@vfamily %in% c("negbinomial", 
        "negbinomial.size")) {
        integer_expression <- TRUE
    }
    else {
        integer_expression <- FALSE
    }
    if (integer_expression) {
        CM <- exprs(cds_subset)
        if (relative_expr) {
            if (is.null(sizeFactors(cds_subset))) {
                stop("Error: to call this function with relative_expr=TRUE, you must call estimateSizeFactors() first")
            }
            CM <- Matrix::t(Matrix::t(CM)/sizeFactors(cds_subset))
        }
        cds_exprs <- reshape2::melt(round(as.matrix(CM)))
    }
    else {
        cds_exprs <- reshape2::melt(exprs(cds_subset))
    }
    if (is.null(min_expr)) {
        min_expr <- cds_subset@lowerDetectionLimit
    }
    colnames(cds_exprs) <- c("f_id", "Cell", "expression")
    cds_pData <- pData(cds_subset)
    cds_fData <- fData(cds_subset)
    cds_exprs <- merge(cds_exprs, cds_fData, by.x = "f_id", by.y = "row.names")
    cds_exprs <- merge(cds_exprs, cds_pData, by.x = "Cell", by.y = "row.names")
    if (integer_expression) {
        cds_exprs$adjusted_expression <- round(cds_exprs$expression)
    }
    else {
        cds_exprs$adjusted_expression <- log10(cds_exprs$expression)
    }
    if (label_by_short_name == TRUE) {
        if (is.null(cds_exprs$gene_short_name) == FALSE) {
            cds_exprs$feature_label <- as.character(cds_exprs$gene_short_name)
            cds_exprs$feature_label[is.na(cds_exprs$feature_label)] <- cds_exprs$f_id
        }
        else {
            cds_exprs$feature_label <- cds_exprs$f_id
        }
    }
    else {
        cds_exprs$feature_label <- cds_exprs$f_id
    }
    cds_exprs$feature_label <- as.factor(cds_exprs$feature_label)
    cds_exprs$Branch <- as.factor(cds_exprs$Branch)
    new_data <- data.frame(Pseudotime = pData(cds_subset)$Pseudotime, 
        Branch = pData(cds_subset)$Branch)
    full_model_expectation <- genSmoothCurves(cds_subset, cores = 1, 
        trend_formula = trend_formula, relative_expr = T, new_data = new_data)
    colnames(full_model_expectation) <- colnames(cds_subset)
    cds_exprs$full_model_expectation <- apply(cds_exprs, 1, function(x) full_model_expectation[x[2], 
        x[1]])
    if (!is.null(reducedModelFormulaStr)) {
        reduced_model_expectation <- genSmoothCurves(cds_subset, 
            cores = 1, trend_formula = reducedModelFormulaStr, 
            relative_expr = T, new_data = new_data)
        colnames(reduced_model_expectation) <- colnames(cds_subset)
        cds_exprs$reduced_model_expectation <- apply(cds_exprs, 
            1, function(x) reduced_model_expectation[x[2], x[1]])
    }
    if (method == "loess") 
        cds_exprs$expression <- cds_exprs$expression + cds@lowerDetectionLimit
    if (label_by_short_name == TRUE) {
        if (is.null(cds_exprs$gene_short_name) == FALSE) {
            cds_exprs$feature_label <- as.character(cds_exprs$gene_short_name)
            cds_exprs$feature_label[is.na(cds_exprs$feature_label)] <- cds_exprs$f_id
        }
        else {
            cds_exprs$feature_label <- cds_exprs$f_id
        }
    }
    else {
        cds_exprs$feature_label <- cds_exprs$f_id
    }
    cds_exprs$feature_label <- factor(cds_exprs$feature_label)
    if (is.null(panel_order) == FALSE) {
        cds_exprs$feature_label <- factor(cds_exprs$feature_label, 
            levels = panel_order)
    }
    cds_exprs$expression[is.na(cds_exprs$expression)] <- min_expr
    cds_exprs$expression[cds_exprs$expression < min_expr] <- min_expr
    cds_exprs$full_model_expectation[is.na(cds_exprs$full_model_expectation)] <- min_expr
    cds_exprs$full_model_expectation[cds_exprs$full_model_expectation < 
        min_expr] <- min_expr
    if (!is.null(reducedModelFormulaStr)) {
        cds_exprs$reduced_model_expectation[is.na(cds_exprs$reduced_model_expectation)] <- min_expr
        cds_exprs$reduced_model_expectation[cds_exprs$reduced_model_expectation < 
            min_expr] <- min_expr
    }
    result <- list(data = cds_exprs, min_expr = min_expr)
    return(result)
}

run_plot <- function(cds_exprs = NULL,
                     min_expr = NULL,
                     color_by = "State",
                     cell_size = 0.75,
                     method = "fitting",
                     nrow = NULL,
                     ncol = 1,
                     reducedModelFormulaStr = NULL,
                     genes = NULL,
                     name = NULL,
                     branch_point = 1){
  raw_data <- cds_exprs
  for(gene in genes){
    cds_exprs <- raw_data %>%  filter(gene_id == gene)
    cds_exprs$State <- as.factor(cds_exprs$State)
    cds_exprs$Branch <- as.factor(cds_exprs$Branch)
    q <- ggplot(aes(Pseudotime, expression), data = cds_exprs)
    if (is.null(color_by) == FALSE) {
      q <- q + geom_point(aes_string(color = color_by), size = I(cell_size))
    }
    if (is.null(reducedModelFormulaStr) == FALSE) 
      q <- q + scale_y_log10()
    else q <- q + scale_y_log10()
    if (method == "loess") 
      q <- q + stat_smooth(aes(fill = Branch, color = Branch), 
                           method = "loess")
    else if (method == "fitting") {
      q <- q + geom_line(aes_string(x = "Pseudotime", y = "full_model_expectation", 
                                    linetype = "Branch"), data = cds_exprs)
    }
    if (!is.null(reducedModelFormulaStr)) {
      q <- q + geom_line(aes_string(x = "Pseudotime", y = "reduced_model_expectation"), 
                         color = "black", linetype = 2, data = cds_exprs)
    }
    q <- q + ylab("Expression") + xlab("Pseudotime (stretched)") + ggtitle(gene)
    q <- q + monocle_theme_opts()
    q <- q + expand_limits(y = min_expr) + scale_color_manual(values = col)
    ggsave(q, filename =paste0(name,"_branch",branch_point,"_",gene,"_pseudotime.png"), width = 6, height = 5, dpi = 300)
    ggsave(q, filename =paste0(name,"_branch",branch_point,"_",gene,"_pseudotime.pdf"), width = 6, height = 5)
  }
}

################################################# 轨迹图上加饼图 #################################################
# blank_theme, 绘制饼图使用, 把各种坐标轴之类的全部去掉
blank_theme <- theme_minimal()+
  theme(
    axis.title.x = element_blank(),
    axis.title.y = element_blank(),
    axis.text.x = element_blank(),
    axis.text.y = element_blank(),
    panel.border = element_blank(),
    panel.grid=element_blank(),
    axis.ticks = element_blank(),
    plot.title=element_text(size=14, face="bold")
  )

# plot pie
# monocl_obj 完成 monocle 分析后的 monocle 对象, 对 monocle 对象使用 pData() 进行数据提取, 过滤出来需要展示的数据并分组统计, 用于绘图
# 例如：想要绘制 Cluster2 和 Cluster 17 在 State1-3 的细胞数占比饼图, 则分组为 Cluster, filter_by 为 State, retain_by 依次为1,2,3 还可以使用自定义的颜色(col)
# plot_pie(monocl_obj = monocl_data, filter_by = "State", retain_by = "1", group_by = "Cluster")
# plot_pie(monocl_obj = monocl_data, filter_by = "State", retain_by = "2", group_by = "Cluster")
# plot_pie(monocl_obj = monocl_data, filter_by = "State", retain_by = "3", group_by = "Cluster")
plot_pie <- function(monocl_obj = NULL,
                     filter_by = NULL,
                     retain_by = NULL,
                     group_by = NULL,
                     col = NULL){
  data <- pData(monocl_obj) %>% 
    tibble::rownames_to_column(var = "cell") %>% 
    filter(.data[[filter_by]] == retain_by) %>% group_by(.data[[group_by]]) %>% count()
  colnames(data) <- c("group","count")
  if(is.null(col)){
    col <- get_colors(colors_len = length(unique(data$group)))
  }
  p <- ggplot2::ggplot(data = data, ggplot2::aes(x = "", y = count, fill = group)) + 
    ggplot2::geom_bar(stat = "identity") +
    coord_polar("y",start = 0) +
    blank_theme +
    theme(legend.position = "none") +
    scale_fill_manual(breaks = unique(data$group),
                      values = col)
  return(p)
}

# 对图进行组合, 需要提供 main_fig 为轨迹, anno_fig 为饼图, x,y 制定了饼图的圆心位置, d 为饼图的直径, 一次只能组合一张饼图
# 组合图片的过程可能要反复进行调整圆心位置 和 直径, 同一张图应该使用同一个 d 值
merge_fig <- function(main_fig = NULL,
                      anno_fig = NULL,
                      x = NULL,
                      y = NULL,
                      d = NULL){
  anno_fig <- ggplotGrob(anno_fig)
  p <- main_fig + annotation_custom(anno_fig, xmin = x - d/2, xmax = x + d/2, 
                                    ymin = y - d/2, ymax = y + d/2)
  return(p)
}